import wfdb
import matplotlib.pyplot as plt

from myutils.NcUtils import fliter
from myutils.beatsLenUtils import extract_fix_length_beats, extract_variable_length_beats, normalize_beat_length

# 使用示例
record = wfdb.rdrecord(f'../data/mit-bih-arrhythmia-database-1.0.0/100', sampfrom=0, sampto=650000, channels=[0])
signal_annotation = wfdb.rdann(f'../data/mit-bih-arrhythmia-database-1.0.0/100', 'atr', sampfrom=0, sampto=650000)
ecg_signal = record.p_signal[:, 0]
r_peaks = signal_annotation.sample
ecg_signal = fliter(ecg_signal)

fix_length_beats, _ = extract_fix_length_beats(ecg_signal, r_peaks)
# 1. 提取变长心拍
variable_length_beats, intervals, _ = extract_variable_length_beats(ecg_signal, r_peaks, signal_annotation.symbol)
print(f"提取了 {len(variable_length_beats)} 个心拍")
for beat, interval in zip(variable_length_beats, intervals):
    print(f"心拍长度: {len(beat)},interval: {interval[1] - interval[0]}")

# 2. (可选) 将心拍归一化为相同长度
target_len = 268
normalized_beats = normalize_beat_length(variable_length_beats, target_len)
print(f"归一化后每个心拍长度: {normalized_beats[0].shape}")

# 设置中文字体（例如黑体 SimHei）
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 或者 ['Microsoft YaHei']
# 解决保存图像时负号'-'显示为方块的问题
plt.rcParams['axes.unicode_minus'] = False

# plt.plot(ecg_signal[42:500], label='000')
# plt.title("ECG信号")
# # plt.plot(fix_length_beats[127], label='原始心拍')
# plt.plot(variable_length_beats[127], label='原始心拍')
# plt.plot(normalized_beats[127], label='处理后心拍')
# # plt.title('第一个等长心拍')
# plt.xlabel('采样点')
# plt.ylabel('振幅')
# plt.legend()
# plt.show()

# 创建2x2子图网格
fig, axs = plt.subplots(2, 2, figsize=(12, 8))  # 设置图形大小[1](@ref)

# 要显示的心拍索引（您可以根据需要修改这些索引）
indices = [345, 2000, 567, 1800]  # 假设您有这些索引的数据

# 将2x2的子图数组展平为一维，便于迭代
axs_flat = axs.ravel()

for i, idx in enumerate(indices):
    ax = axs_flat[i]

    # 检查索引是否在数据范围内
    if idx < len(variable_length_beats) and idx < len(normalized_beats):
        # 在当前子图绘制原始心拍和处理后心拍
        ax.plot(variable_length_beats[idx], label='原始心拍', color='blue', alpha=0.7)
        ax.plot(normalized_beats[idx], label='处理后心拍', color='red', alpha=0.8)

        ax.set_title(f'心拍 #{i + 1}对比')  # 设置子图标题[4](@ref)
        ax.set_xlabel('采样点')
        ax.set_ylabel('振幅')
        ax.legend()  # 显示图例
        ax.grid(True, alpha=0.3)  # 添加网格线便于观察
    else:
        # 如果索引超出范围，显示提示信息
        ax.text(0.5, 0.5, f'索引 {idx} 超出数据范围',
                ha='center', va='center', transform=ax.transAxes, fontsize=12)
        ax.set_title(f'心拍 #{idx} (无数据)')

# 自动调整子图间距，避免重叠[1,5](@ref)
plt.tight_layout(pad=2.0)  # 增加内边距使布局更宽松

# 添加整个图形的总标题（可选）
fig.suptitle('ECG心拍信号对比分析', fontsize=16, y=0.98)

# plt.show()
plt.savefig('1.jpg', dpi=600, bbox_inches='tight')
